package main

import (
	"fmt"
	"io"
	"net"
	"time"

	"gitee.com/sjy/studyGo/common"
	"github.com/astaxie/beego/logs"
	"github.com/xtaci/smux"
)

type Client struct {
	clientID   string
	serverAddr string
}

func NewClient(clientID string, serverAddr string) *Client {
	return &Client{
		clientID,
		serverAddr,
	}
}

func (c *Client) Run() {
	for {

		func() {
			defer func() {
				if err := recover(); err != nil && err != io.EOF {
					logs.Error("%v", err)
				}
			}()
			// 安全的调用 run
			c.run()
		}()

		logs.Warn("reconnect %s", c.serverAddr)
		time.Sleep(time.Second * 1)
	}
}

func (c *Client) run() {
	conn, err := net.Dial("tcp", c.serverAddr)
	if err != nil {
		panic(err)
	}
	defer conn.Close()

	// 发送handshake包
	handshakeReq := common.HandshakeReq{ClientID: c.clientID}
	buf := handshakeReq.Encode()

	conn.SetWriteDeadline(time.Now().Add(time.Second * 3))
	_, err = conn.Write(buf)
	conn.SetWriteDeadline(time.Time{})
	if err != nil {
		panic(err)
	}

	// 创建mux session
	mux, err := smux.Client(conn, nil)
	if err != nil {
		panic(err)
	}
	defer mux.Close()

	// 等待mux stream
	for {
		stream, err := mux.AcceptStream()
		if err != nil {
			panic(err)
		}

		go c.handleStream(stream)
	}
}

func (c *Client) handleStream(stream net.Conn) {
	defer stream.Close()

	func() {
		defer func() {
			if err := recover(); err != nil {
				logs.Error("%v", err)
			}
		}()
		// pp解码
		pp := &common.ProxyProtocol{}
		pp.Decode(stream)

		// 与本地建连接
		switch pp.InternalProtocol {
		case "tcp":
			handleTcp(pp, stream)
		case "udp":
			handleUdp(pp, stream)
		default:
			logs.Warn("unsupported protocol %s", pp.InternalProtocol)
		}
	}()
}

func handleTcp(pp *common.ProxyProtocol, stream net.Conn) {
	localConn, err := net.Dial("tcp", fmt.Sprintf("%s:%d", pp.InternalIP, pp.InternalPort))
	if err != nil {
		panic(fmt.Errorf("connect to to local fail: %v", err))
	}
	defer localConn.Close()

	// 双向数据拷贝
	go func() {
		defer localConn.Close()
		defer stream.Close()
		io.Copy(localConn, stream)
	}()
	io.Copy(stream, localConn)
}

func handleUdp(pp *common.ProxyProtocol, stream net.Conn) {
	localConn, err := net.Dial("udp", fmt.Sprintf("%s:%d", pp.InternalIP, pp.InternalPort))
	if err != nil {
		panic(fmt.Errorf("connect to to local fail: %v", err))
	}
	defer localConn.Close()

	// read local conn
	go func() {
		defer localConn.Close()
		defer stream.Close()
		buf := make([]byte, 1024*64)
		for {
			nr, err := localConn.Read(buf)
			if err != nil {
				logs.Error("read udp from local fail %v", err)
				break
			}

			logs.Debug("read %d bytes from local connect", nr)
			// udp packet编码
			p := common.UDPPacket(buf[:nr])
			body, err := p.Encode()
			if err != nil {
				logs.Warn("encode udp packet fail: %v", err)
				break
			}

			_, err = stream.Write(body)
			if err != nil {
				logs.Warn("write udp to stream fail: %v", err)
				break
			}
		}
	}()

	// read stream
	p := common.UDPPacket(make([]byte, 1024*64))
	for {
		nr, err := p.Decode(stream)
		if err != nil {
			logs.Warn("decode udp from stream fail: %v", err)
			break
		}

		logs.Debug("read from stream %d bytes", nr)
		_, err = localConn.Write(p[:nr])
		if err != nil {
			logs.Warn("write udp to local conn fail: %v", err)
			break
		}
	}
}
